{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f8dbcd99790>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Copyright (c) Microsoft Corporation.\n",
    "# Licensed under the MIT license.\n",
    "\n",
    "import os\n",
    "import re\n",
    "import json\n",
    "import gzip\n",
    "import torch\n",
    "import pickle\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from collections import defaultdict\n",
    "\n",
    "def load_pickle(filename):\n",
    "    with open(filename, \"rb\") as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "def save_pickle(data, filename):\n",
    "    with open(filename, \"wb\") as f:\n",
    "        pickle.dump(data, f)\n",
    "\n",
    "def load_json(file_path):\n",
    "    with open(file_path, \"r\") as f:\n",
    "        return json.load(f)\n",
    "    \n",
    "def ReadLineFromFile(path):\n",
    "    lines = []\n",
    "    with open(path,'r') as fd:\n",
    "        for line in fd:\n",
    "            lines.append(line.rstrip('\\n'))\n",
    "    return lines\n",
    "\n",
    "def parse(path):\n",
    "    g = gzip.open(path, 'r')\n",
    "    for l in g:\n",
    "        yield eval(l)\n",
    "        \n",
    "'''\n",
    "Set seeds\n",
    "'''\n",
    "seed = 2022\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "short_data_name = 'ml1m'\n",
    "full_data_name = 'ml-1m'\n",
    "if not os.path.exists(os.path.join(\"../../data/\", short_data_name)):\n",
    "    os.mkdir(os.path.join(\"../../data/\", short_data_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# return (user, item, timestamp) sort in get_interaction\n",
    "def ML1M(rating_score=3):\n",
    "    datas = []\n",
    "    data_dict = {}\n",
    "    data_file = \"../../data/raw_data/ml-1m/ratings.dat\"\n",
    "\n",
    "    inter_df = pd.read_csv(data_file, sep='::', header=None)\n",
    "    inter_df.columns = [\"uid\", \"iid\", \"rating\", \"timestamp\"]\n",
    "    inter_df = inter_df.sort_values(by=\"timestamp\")\n",
    "    inter_df = inter_df.drop_duplicates([\"uid\", \"iid\"]).reset_index(drop=True)\n",
    "\n",
    "    for idx in range(len(inter_df)):\n",
    "        user = str(inter_df.loc[idx, \"uid\"])\n",
    "        item = str(inter_df.loc[idx, \"iid\"])\n",
    "        time = inter_df.loc[idx, \"timestamp\"]\n",
    "        rating = inter_df.loc[idx, \"rating\"]\n",
    "        if rating <= rating_score:\n",
    "            continue\n",
    "        data_dict[(user, item)] = int(time)\n",
    "        datas.append((user, item, int(time)))\n",
    "    return datas\n",
    "\n",
    "def ML1M_meta(datamaps):\n",
    "    meta_datas = {}\n",
    "    genres, years = defaultdict(int), defaultdict(int)\n",
    "    meta_file = \"../../data/raw_data/ml-1m/movies.dat\"\n",
    "    item_ids = list(datamaps['item2id'].keys())\n",
    "    item_df = pd.read_csv(meta_file, sep='::', header=None, encoding=\"ISO-8859-1\")\n",
    "    item_df.columns = [\"iid\", 'i_title', 'i_genre']\n",
    "    \n",
    "    for idx in range(len(item_df)):\n",
    "        iid = str(item_df.loc[idx, \"iid\"])\n",
    "        if iid not in item_ids:\n",
    "            continue\n",
    "        title = item_df.loc[idx, \"i_title\"]\n",
    "        item_title = title[:-7]\n",
    "        item_year = title[-5:-1]\n",
    "        item_genre = item_df.loc[idx, \"i_genre\"].replace(\"|\", \", \")\n",
    "        meta = {\"id\": str(iid), \"title\": item_title, \"year\": item_year, \"genre\": item_genre}\n",
    "        mapped_id = datamaps['item2id'][iid]\n",
    "        meta_datas[mapped_id] = meta\n",
    "\n",
    "    user_datas = {}\n",
    "    user_file = \"../../data/raw_data/ml-1m/users.dat\"\n",
    "    user_ids = list(datamaps['user2id'].keys())\n",
    "    user_df = pd.read_csv(user_file, sep='::', header=None)\n",
    "    user_df.columns = [\"uid\", \"gender\", \"age\", \"occupation\", \"zip_code\"]\n",
    "    for idx in range(len(user_df)):\n",
    "        uid = str(user_df.loc[idx, \"uid\"])\n",
    "        if uid not in user_ids:\n",
    "            continue\n",
    "        gender = user_df.loc[idx, \"gender\"]\n",
    "        age = user_df.loc[idx, \"age\"]\n",
    "        occupation = user_df.loc[idx, \"occupation\"]\n",
    "        zip_code = user_df.loc[idx, \"zip_code\"]\n",
    "        user = {\"id\": str(uid), \"gender\": gender, \"age\": age, \"occupation\": occupation, \"zip_code\": zip_code}\n",
    "\n",
    "    datamaps['genres'] = genres\n",
    "    return meta_datas, datamaps, user_datas\n",
    "     \n",
    "def add_comma(num): # 1000000 -> 1,000,000\n",
    "    str_num = str(num)\n",
    "    res_num = ''\n",
    "    for i in range(len(str_num)):\n",
    "        res_num += str_num[i]\n",
    "        if (len(str_num)-i-1) % 3 == 0:\n",
    "            res_num += ','\n",
    "    return res_num[:-1]\n",
    "\n",
    "# get user interaction sequence for sequential recommendation\n",
    "def get_interaction(datas):\n",
    "    user_seq = {}\n",
    "    for data in datas:\n",
    "        user, item, time = data\n",
    "        if user in user_seq:\n",
    "            user_seq[user].append((item, time))\n",
    "        else:\n",
    "            user_seq[user] = []\n",
    "            user_seq[user].append((item, time))\n",
    "\n",
    "    for user, item_time in user_seq.items():\n",
    "        item_time.sort(key=lambda x: x[1])  \n",
    "        items = []\n",
    "        for t in item_time:\n",
    "            items.append(t[0])\n",
    "        user_seq[user] = items\n",
    "    return user_seq\n",
    "\n",
    "# K-core user_core item_core, return False if any user/item < core\n",
    "def check_Kcore(user_items, user_core, item_core):\n",
    "    user_count = defaultdict(int)\n",
    "    item_count = defaultdict(int)\n",
    "    for user, items in user_items.items():\n",
    "        for item in items:\n",
    "            user_count[user] += 1\n",
    "            item_count[item] += 1\n",
    "\n",
    "    for user, num in user_count.items():\n",
    "        if num < user_core:\n",
    "            return user_count, item_count, False\n",
    "    for item, num in item_count.items():\n",
    "        if num < item_core:\n",
    "            return user_count, item_count, False\n",
    "    return user_count, item_count, True # Kcore guaranteed\n",
    "\n",
    "# recursively K-core filtering \n",
    "def filter_Kcore(user_items, user_core, item_core):  \n",
    "    user_count, item_count, isKcore = check_Kcore(user_items, user_core, item_core)\n",
    "    while not isKcore:\n",
    "        for user, num in user_count.items():\n",
    "            if user_count[user] < user_core: # remove the user\n",
    "                user_items.pop(user)\n",
    "            else:\n",
    "                for item in user_items[user]:\n",
    "                    if item_count[item] < item_core:\n",
    "                        user_items[user].remove(item)\n",
    "        user_count, item_count, isKcore = check_Kcore(user_items, user_core, item_core)\n",
    "    return user_items\n",
    "\n",
    "def id_map(user_items): # user_items dict\n",
    "    user2id = {} # raw 2 uid\n",
    "    item2id = {} # raw 2 iid\n",
    "    id2user = {} # uid 2 raw\n",
    "    id2item = {} # iid 2 raw\n",
    "    user_id = 1  # start from 1\n",
    "    item_id = 1\n",
    "    final_data = {}\n",
    "    random_user_list = list(user_items.keys())\n",
    "    random.shuffle(random_user_list)  # shuffle for re-indexing\n",
    "    for user in random_user_list:\n",
    "        items = user_items[user]\n",
    "        if user not in user2id:\n",
    "            user2id[user] = str(user_id)\n",
    "            id2user[str(user_id)] = user\n",
    "            user_id += 1\n",
    "        iids = [] # item id lists\n",
    "        for item in items:\n",
    "            if item not in item2id:\n",
    "                item2id[item] = str(item_id)\n",
    "                id2item[str(item_id)] = item\n",
    "                item_id += 1\n",
    "            iids.append(item2id[item])\n",
    "        uid = user2id[user]\n",
    "        final_data[uid] = iids\n",
    "    data_maps = {\n",
    "        'user2id': user2id,\n",
    "        'item2id': item2id,\n",
    "        'id2user': id2user,\n",
    "        'id2item': id2item\n",
    "    }\n",
    "    return final_data, user_id-1, item_id-1, data_maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(data_name, acronym, data_type='ML1M'):\n",
    "    assert data_type in {'Amazon', 'Yelp', 'Steam', 'ML100k', 'ML1M'}\n",
    "    rating_score = 3.0  # rating score smaller than this score would be deleted\n",
    "    # user 5-core item 5-core\n",
    "    user_core = 5\n",
    "    item_core = 1\n",
    "    attribute_core = 0\n",
    "\n",
    "    datas = ML1M(rating_score)  # list of [user, item, timestamp]\n",
    "\n",
    "    user_items = get_interaction(datas) # dict of {user: interaction list sorted by time} \n",
    "    print(f'{data_name} Raw data has been processed! Lower than {rating_score} are deleted!')\n",
    "    # raw_id user: [item1, item2, item3...]\n",
    "    user_items = filter_Kcore(user_items, user_core=user_core, item_core=item_core)\n",
    "    print(f'User {user_core}-core complete! Item {item_core}-core complete!')\n",
    "\n",
    "    user_items, user_num, item_num, data_maps = id_map(user_items) # get mapping dicts, randomly shuffle\n",
    "    user_count, item_count, _ = check_Kcore(user_items, user_core=user_core, item_core=item_core)\n",
    "    user_count_list = list(user_count.values()) # user click count\n",
    "    user_avg, user_min, user_max = np.mean(user_count_list), np.min(user_count_list), np.max(user_count_list)\n",
    "    item_count_list = list(item_count.values()) # item click count\n",
    "    item_avg, item_min, item_max = np.mean(item_count_list), np.min(item_count_list), np.max(item_count_list)\n",
    "    interact_num = np.sum([x for x in user_count_list])\n",
    "    sparsity = (1 - interact_num / (user_num * item_num)) * 100\n",
    "    show_info = f'Total User: {user_num}, Avg User: {user_avg:.4f}, Min Len: {user_min}, Max Len: {user_max}\\n' + \\\n",
    "                f'Total Item: {item_num}, Avg Item: {item_avg:.4f}, Min Inter: {item_min}, Max Inter: {item_max}\\n' + \\\n",
    "                f'Iteraction Num: {interact_num}, Sparsity: {sparsity:.2f}%'\n",
    "    print(show_info)\n",
    "\n",
    "\n",
    "    print('Begin extracting meta infos...')\n",
    "    \n",
    "    meta_infos, datamaps, user_infos = ML1M_meta(data_maps)\n",
    "\n",
    "    print(f'{data_name} & {add_comma(user_num)} & {add_comma(item_num)} & {user_avg:.1f}'\n",
    "          f'& {item_avg:.1f} & {add_comma(interact_num)} & {sparsity:.2f}\\% \\\\')\n",
    "\n",
    "    # -------------- Save Data ---------------\n",
    "    data_file = '../../data/{}/'.format(acronym) + 'sequential_data.txt'\n",
    "    metadata_file = '../../data/{}/'.format(acronym) + 'metadata.json'\n",
    "    userdata_file = '../../data/{}/'.format(acronym) + 'userdata.json'\n",
    "    datamaps_file = '../../data/{}/'.format(acronym) + 'datamaps.json'\n",
    "\n",
    "    with open(data_file, 'w') as out:\n",
    "        for user, items in user_items.items():\n",
    "            out.write(user + ' ' + ' '.join(items) + '\\n')\n",
    "\n",
    "    item_keys = sorted(meta_infos.keys(), key=lambda x: int(x))\n",
    "    print(f\"item2id: {len(datamaps['item2id'])}, meta_infos: {len(meta_infos)}, item_keys: {item_keys[:100]}\")\n",
    "    with open(metadata_file, 'w') as out:\n",
    "        for key in item_keys:\n",
    "            out.write(json.dumps(meta_infos[key]) + '\\n')\n",
    "\n",
    "    # print(\"datamap: \", datamaps.keys(), datamaps['item2id'].keys(), datamaps['user2id'].keys(), datamaps['id2item'].keys(), datamaps['id2user'].keys())\n",
    "    json_str = json.dumps(datamaps)\n",
    "    with open(datamaps_file, 'w') as out:\n",
    "        out.write(json_str)\n",
    "\n",
    "    # -------------- Split Train/Valid/Test for Item Import & Tagging ---------------\n",
    "    all_items = [item for item in datamaps['item2id'].keys()]\n",
    "    random.shuffle(all_items)\n",
    "    train_split = int(len(all_items) * 0.8)\n",
    "    valid_split = int(len(all_items) * 0.1)\n",
    "    train_items = all_items[:train_split]\n",
    "    valid_items = all_items[train_split:train_split+valid_split]\n",
    "    test_items = all_items[train_split+valid_split:]\n",
    "    outputs = {'train': train_items, 'val': valid_items, 'test': test_items}\n",
    "    save_pickle(outputs, '../../data/{}/item_splits.pkl'.format(short_data_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_593252/807186883.py:7: ParserWarning: Falling back to the 'python' engine because the 'c' engine does not support regex separators (separators > 1 char and different from '\\s+' are interpreted as regex); you can avoid this warning by specifying engine='python'.\n",
      "  inter_df = pd.read_csv(data_file, sep='::', header=None)\n"
     ]
    }
   ],
   "source": [
    "main(full_data_name, short_data_name, data_type='ML1M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_test_data(data_name, test_num=99, sample_type='random'):\n",
    "    \"\"\"\n",
    "    sample_type:\n",
    "        random:  sample `test_num` negative items randomly.\n",
    "        pop: sample `test_num` negative items according to item popularity.\n",
    "    \"\"\"\n",
    "\n",
    "    data_file = f'sequential_data.txt'\n",
    "    if sample_type == 'pop':\n",
    "        test_file = f'negative_samples_pop.txt'\n",
    "    else:\n",
    "        test_file = f'negative_samples.txt'\n",
    "\n",
    "    item_count = defaultdict(int)\n",
    "    user_items = defaultdict()\n",
    "\n",
    "    lines = open('../../data/{}/'.format(data_name) + data_file).readlines()\n",
    "    for line in tqdm(lines):\n",
    "        user, items = line.strip().split(' ', 1)\n",
    "        items = items.split(' ')\n",
    "        items = [int(item) for item in items]\n",
    "        user_items[user] = items\n",
    "        for item in items:\n",
    "            item_count[item] += 1\n",
    "\n",
    "    all_item = list(item_count.keys())\n",
    "    count = list(item_count.values())\n",
    "    sum_value = np.sum([x for x in count])\n",
    "    probability = [value / sum_value for value in count]\n",
    "\n",
    "    user_neg_items = defaultdict()\n",
    "\n",
    "    for user, user_seq in tqdm(user_items.items()):\n",
    "        test_samples = []\n",
    "        while len(test_samples) < test_num:\n",
    "            if sample_type == 'random':\n",
    "                sample_ids = np.random.choice(all_item, test_num, replace=False)\n",
    "            else: # sample_type == 'pop':\n",
    "                np.random.seed(int(random.random() * 1000))\n",
    "                sample_ids = np.random.choice(all_item, test_num, replace=False, p=probability)\n",
    "            sample_ids = [str(item) for item in sample_ids if item not in user_seq and item not in test_samples]\n",
    "            test_samples.extend(sample_ids)\n",
    "        test_samples = test_samples[:test_num]\n",
    "        user_neg_items[user] = test_samples\n",
    "\n",
    "    with open('../../data/{}/'.format(data_name) + test_file, 'w') as out:\n",
    "        for user, samples in user_neg_items.items():\n",
    "            out.write(user+' '+' '.join(samples)+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6034/6034 [00:00<00:00, 43414.48it/s]\n",
      "  0%|          | 0/6034 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6034/6034 [00:57<00:00, 105.63it/s]\n"
     ]
    }
   ],
   "source": [
    "sample_test_data(short_data_name, test_num=99, sample_type=\"random\")  #sample 99 negative testing samples for each user\n",
    "sample_test_data(short_data_name, test_num=99, sample_type=\"pop\")  #sample 99 negative testing samples for each user"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "183e66c1b4a3cebcc9a59d658c813a22bf955a401fdefcad593bde0024c573fe"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
